import pathlib
import matplotlib.pyplot as plt
from tensorflow import keras
import tensorflow as tf

plt.rcParams['font.sans-serif'] = ['SimHei']

path = "../../static/img"
# 解析目录
data_dir = pathlib.Path(path)

# keras 加载数据集
batch_size = 32
img_height = 180
img_width = 180

# 使用 80% 的图像进行训练，20% 的图像进行验证。
class_names = ['Battery', 'BrickAndTileCeramics', 'Cans', 'cigarette',
               'Fruits', 'NO_RUBBISH', 'Vegetables', 'WaterBottle']
class_names_cn = ['电池', '砖瓦陶瓷', '罐', '香烟',
                  '水果', '没有垃圾', '蔬菜', '水壶']

train_ds = keras.utils.image_dataset_from_directory(
    directory=data_dir,
    validation_split=0.2,
    subset="training",
    image_size=(img_height, img_width),
    batch_size=batch_size,
    shuffle=True,
    seed=123,
    interpolation='bilinear',
    crop_to_aspect_ratio=True,
    labels='inferred',
    class_names=class_names,
    color_mode='rgb',
)

AUTOTUNE = tf.data.AUTOTUNE
train_ds = train_ds.cache().prefetch(buffer_size=AUTOTUNE)

model = keras.Sequential([
    # 归一化
    keras.layers.Rescaling(1. / 255),
    # 随机旋转
    keras.layers.RandomRotation(0.2)
])

# 获取第一个批次数据
for image_batch, labels_batch in train_ds.take(1):
    for i in range(len(image_batch)):
        plt.figure(figsize=(10, 10))
        for j in range(9):
            plt.subplot(3, 3, j + 1)
            augmented_image = model.predict(tf.expand_dims(image_batch[i], 0))
            print(augmented_image[0].shape)
            plt.imshow(augmented_image[0])
            plt.title(class_names_cn[labels_batch[i]])
            plt.axis('off')
        plt.show()
